import cv2
import numpy as np
import glob

# 棋盘格参数
chessboard_size = (9, 6)  # 棋盘格内角点的数量
square_size = 1.0  # 棋盘格每个方格的边长，单位可以是毫米、厘米等

# 准备对象点，如 (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2) * square_size

# 存储对象点和图像点的列表
objpoints = []  # 3d点在世界坐标系中的位置
imgpoints = []  # 2d点在图像平面中的位置


# 读取所有标定图像
images = glob.glob('calibration_images/*.jpg')

for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 查找棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

    if ret:
        objpoints.append(objp)
        imgpoints.append(corners)

        # 绘制并显示角点
        cv2.drawChessboardCorners(img, chessboard_size, corners, ret)
        cv2.imshow('img', img)
        cv2.waitKey(500)

cv2.destroyAllWindows()

# 进行相机标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)


# 计算重投影误差
mean_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
    mean_error += error

print(f"Total re-projection error: {mean_error / len(objpoints)}")

# 保存标定结果
np.savez('calibration_output.npz', ret=ret, mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)

# 读取标定结果
calibration_data = np.load('calibration_output.npz')
ret = calibration_data['ret']
mtx = calibration_data['mtx']
dist = calibration_data['dist']
rvecs = calibration_data['rvecs']
tvecs = calibration_data['tvecs']

# 打印标定结果
print("Camera matrix:\n", mtx)
print("Distortion coefficients:\n", dist)
print("Rotation vectors:\n", rvecs)
print("Translation vectors:\n", tvecs)


# 读取一张图像并校正畸变
img = cv2.imread('calibration_images/1.jpg')
h, w = img.shape[:2]
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))

# 校正图像
dst = cv2.undistort(img, mtx, dist, None, newcameramtx)

# 剪裁图像
x, y, w, h = roi
dst = dst[y:y+h, x:x+w]
cv2.imwrite('calibrated_image.png', dst)